#include "kernel_operator.h"
using namespace AscendC;

class NLLLoss1d {
 public:
  __aicore__ inline NLLLoss1d() {}
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight,
                              GM_ADDR y, uint32_t sizeC,
                              uint32_t reduction, int32_t ignoreIndex) {
    // printf("------------dimNum==1------------\n");
    this->reduction = reduction;
    this->ignoreIndex = ignoreIndex;

    xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x, sizeC * sizeof(DTYPE_X));
    targetGm.SetGlobalBuffer((__gm__ DTYPE_TARGET*)target,
                             sizeof(DTYPE_TARGET));
    weightGm.SetGlobalBuffer((__gm__ DTYPE_WEIGHT*)weight,
                             sizeC * sizeof(DTYPE_WEIGHT));
    yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, sizeof(DTYPE_Y));
  }

  __aicore__ inline void process() {
    int target = targetGm.GetValue(0);

    if (target == this->ignoreIndex) {
      yGm.SetValue(0, 0);
      return;
    }
    float x = xGm.GetValue(target);
    if (this->reduction == 2) {  // "mean"
      yGm.SetValue(0, x * -1);
    } else {
      float w = weightGm.GetValue(target);
      yGm.SetValue(0, x * w * -1);
    }
    return;
  }

 private:
  uint32_t reduction;
  int32_t ignoreIndex;

  // AscendC::TPipe pipe;
  AscendC::GlobalTensor<DTYPE_X> xGm;
  AscendC::GlobalTensor<DTYPE_TARGET> targetGm;
  AscendC::GlobalTensor<DTYPE_WEIGHT> weightGm;
  AscendC::GlobalTensor<DTYPE_Y> yGm;
};